from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import warnings
warnings.filterwarnings("ignore")

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from options import Options
from lib.data import Dataset
from lib.timegan import TimeGAN


def impute(model_dir):
    """ impute
    """
    # ARGUMENTS
    opt = Options().parse()
    opt.resume = True  # resume = True 初始化时会先加载，不用再训练了
    opt.model_dir = model_dir

    # LOAD DATA
    dataset = Dataset(opt.data_name, opt.data_type, opt.seq_len, opt.indicators, opt.masked_indicator)
    masked_data = dataset.load_masked_data()
    # LOAD MODEL
    model = TimeGAN(opt, dataset)
    # # 先训练，如果训练了直接加载
    # model.train()
    # impute MODEL
    model.imputation(masked_data)

if __name__ == '__main__':
    impute()